#coding:utf-8
'''
训练模型 + 评估
1. load_dataset 数据读入：快速训练-采样
2. 数据处理：
    (1) train: map预处理: tokenizer + remove_columns
    (2) test: 数据加载器data_loader
3. model 模型定义
4. train 模型训练
5. 测试评估
'''
import torch
import numpy as np
from datasets import load_from_disk
from datasets import load_metric
from transformers.trainer_utils import EvalPrediction
from transformers import TrainingArguments, Trainer

from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

'''
1. load_dataset 数据读入：快速训练-采样
'''
#从本地磁盘加载数据
datasets = load_from_disk('./data/ChnSentiCorp')
#取数据子集，否则数据太多跑不动
dataset_train = datasets['train'].shuffle().select(range(1000))       #打乱+选取数据: 1000
dataset_test = datasets['validation'].shuffle().select(range(200))    #200
del datasets
# 不取子集
#dataset_train = datasets['train'].shuffle()
#dataset_test = datasets['validation'].shuffle()

'''
2. 数据处理
'''
'''
train: map预处理: tokenizer + remove_columns
'''
#加载分词工具
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')   #---选定分词器

#针对每个数据进行分词
def f(data):
    return tokenizer(
        data['text'],         #---查看前面数据的内容
        padding='max_length', #一律补零到max_length长度
        truncation=True,      #超过就截断
        max_length=30,
    )
#批量处理
dataset_train = dataset_train.map(f, batched=True, batch_size=1000, num_proc=4)
dataset_test = dataset_test.map(f, batched=True, batch_size=100, num_proc=4)

'''
3. model 模型定义
'''
#加载模型
model = AutoModelForSequenceClassification.from_pretrained('bert-base-cased',  #---模型
                                                           num_labels=2)       #---类别
#加载评价函数
metric = load_metric('accuracy')
#定义评价函数
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    logits = logits.argmax(axis=1)
    return metric.compute(predictions=logits, references=labels)

'''
4. train 模型训练
'''
#初始化训练参数
args = TrainingArguments(output_dir='./output_dir', evaluation_strategy='epoch')
args.num_train_epochs = 1
args.learning_rate = 1e-4
args.weight_decay = 1e-2
args.per_device_eval_batch_size = 32
args.per_device_train_batch_size = 16

#初始化训练器
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=dataset_train,        # 训练集
    eval_dataset=dataset_test,          # 测试集
    compute_metrics=compute_metrics,    # 评价函数
)

#评价模型
res = trainer.evaluate()
print('评价函数：')
print(res)
#保存模型
trainer.save_model(output_dir='./output_dir')
#config.json：参数配置
#pytorch_model.bin：模型
#training_args.bin：模型参数

